import os

import albumentations as A
import cv2
import numpy as np
import pandas as pd
import torch
from pandas import DataFrame
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torch.utils.data import Dataset

from preprocessing.constants import CASIA_FASD
from timm.datasets.transform import create_train_transform, create_val_test_transform


class CasiaFASDDataset(Dataset):

    def __init__(self, data_root, df: DataFrame, mode, transform: A.Compose, color_space='rgb'):
        """Initialization"""
        self.data_root = data_root
        self.df = df
        self.mode = mode
        self.transform = transform
        self.color_space = color_space

    def __getitem__(self, index):
        video, img_file, label = self.df.iloc[index].values
        if self.mode == 'train':
            rgb_path = os.path.join(self.data_root, 'train_release_crops', video, img_file)
            msr_path = os.path.join(self.data_root, 'train_release_MSRCR', video, img_file)
        elif self.mode == 'validation':
            rgb_path = os.path.join(self.data_root, 'train_release_crops', video, img_file)
            msr_path = os.path.join(self.data_root, 'train_release_MSRCR', video, img_file)

        else:
            rgb_path = os.path.join(self.data_root, 'test_release_crops', video, img_file)
            msr_path = os.path.join(self.data_root, 'test_release_MSRCR', video, img_file)



        msr = cv2.imread(msr_path, cv2.IMREAD_GRAYSCALE)
        image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
        if self.color_space == 'rgb':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image, msr = self._get_transformed(image, msr)  # RGB
        elif self.color_space == 'hsv':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
            image, msr = self._get_transformed(image, msr)  # HSV
        elif self.color_space == 'ycrcb':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
            image, msr = self._get_transformed(image, msr)  # YCrCb
        elif self.color_space == 'full':
            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
            ycrcb = cv2.cvtColor(image, cv2.COLOR_BGR2YCrCb)
            sample = {'image': rgb, 'image2': hsv, 'image3': ycrcb, 'mask': msr}
            transformed = self.transform(**sample)
            rgb, hsv, ycrcb = transformed["image"], transformed["image2"], transformed["image3"]
            msr = transformed["mask"] / 255.
            msr = msr.unsqueeze(0).repeat(3, 1, 1)
            image = np.concatenate((rgb, hsv, ycrcb), axis=0)
        else:
            raise NotImplementedError("Color space only include rgb, hsv, ycrcb, full!")

        return {'images': image, 'msrs': msr, 'labels': label}

    def _get_transformed(self, image, msr):
        transformed = self.transform(image=image, mask=msr)
        image = transformed["image"]
        msr = transformed["mask"] / 255.
        msr = msr.unsqueeze(0).repeat(3, 1, 1)
        return image, msr

    def __len__(self):
        r = self.df.shape[0]
        return r


def get_casia_fasd_dataloader(model_cfg, args, color_space):
    df = pd.read_csv(f'data/{CASIA_FASD}/data_train.csv')
    train_df, val_df = train_test_split(df, test_size=0.1, random_state=111)
    # train_df = train_df.iloc[:1000]
    train_transform = create_train_transform(model_cfg)
    train_data = CasiaFASDDataset(data_root=args.data_dir, df=train_df, mode='train', transform=train_transform,
                                  color_space=color_space)
    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
    else:
        train_sampler = None
    train_loader = DataLoader(train_data, args.batch_size, shuffle=(train_sampler is None),
                              sampler=train_sampler, num_workers=args.workers, pin_memory=True, drop_last=True)

    # val_df = val_df.iloc[:1000]
    val_transform = create_val_test_transform(model_cfg)
    val_data = CasiaFASDDataset(data_root=args.data_dir, df=val_df, mode='validation', transform=val_transform,
                                color_space=color_space)
    val_loader = DataLoader(val_data, 30, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=False)
    return train_sampler, train_loader, val_loader


def get_casia_fasd_test_dataloader(model_cfg, args, color_space):
    test_df = pd.read_csv(f'data/{args.prefix}/data_test.csv')
    # test_df = test_df.iloc[:2000]
    test_transform = create_val_test_transform(model_cfg)
    test_data = CasiaFASDDataset(data_root=args.data_dir, df=test_df, mode='test', transform=test_transform,
                                 color_space=color_space)
    test_loader = DataLoader(test_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers,
                             pin_memory=True, drop_last=False)
    return test_loader
